#!/usr/bin/env python3
"""
PyInstaller 打包脚本
"""
import os
import shutil
import subprocess
import sys
from pathlib import Path

def clean_build_dirs():
    """清理构建目录"""
    dirs_to_clean = ['build', 'dist']
    for dir_name in dirs_to_clean:
        if Path(dir_name).exists():
            print(f"清理目录: {dir_name}")
            shutil.rmtree(dir_name)

def install_dependencies():
    """安装依赖"""
    print("安装依赖包...")
    requirements = [
        'rembg',
        'tqdm',
        'pillow',
        'numpy',
        'onnxruntime',
        'requests',
        'pooch',
        'scipy',
        'pymatting',
    ]
    
    for package in requirements:
        print(f"安装 {package}...")
        subprocess.check_call([sys.executable, '-m', 'pip', 'install', package])

def build_executable():
    """构建可执行文件"""
    print("开始构建可执行文件...")
    
    # 使用 spec 文件构建
    result = subprocess.run([
        'pyinstaller',
        '--clean',
        'rmbg.spec'
    ], capture_output=True, text=True)
    
    if result.returncode != 0:
        print("构建失败!")
        print("STDOUT:", result.stdout)
        print("STDERR:", result.stderr)
        return False
    
    print("构建成功!")
    return True

def create_models_directory():
    """创建模型目录"""
    models_dir = Path('dist') / 'models'
    models_dir.mkdir(parents=True, exist_ok=True)
    print(f"创建模型目录: {models_dir}")

def main():
    """主函数"""
    print("=== BG Remover 打包工具 ===")
    
    try:
        # 清理旧构建
        clean_build_dirs()
        
        # 安装依赖
        # install_dependencies()
        
        # 构建可执行文件
        if build_executable():
            # 创建模型目录
            create_models_directory()
            
            print("\n=== 打包完成 ===")
            print("可执行文件位置: dist/rmbg/")
            print("模型目录: dist/rmbg/models/")
            print("\n使用说明:")
            print("1. 首次运行会自动下载模型文件到 models 目录")
            print("2. 可以通过 --model-dir 参数指定其他模型目录")
        else:
            print("打包失败!")
            sys.exit(1)
            
    except Exception as e:
        print(f"打包过程中出现错误: {e}")
        sys.exit(1)

if __name__ == "__main__":
    main()